
#############################################################
# 本脚本用于截取 OBS 中指定场景的截图，并保存到指定路径
# 请确保已经安装 simpleobsws 库
# pip install simpleobsws
# 请确保已经在 OBS 中启用了 WebSocket 服
# 请修改 main 函数中的 scenes 列表为你实际的场景名称
# 请修改 take_screenshot 函数中的 screenshot_path 为你实际的保存路径
# 请修改 take_screenshot 函数中的 sourceName 为你实际的场景名称
# KZENG 2021-09-07
#############################################################


import asyncio
import simpleobsws
import time

# 设置连接信息
# parameters = simpleobsws.IdentificationParameters() # 默认连接到 ws://localhost:4455
# ws = simpleobsws.WebSocketClient(parameters)

parameters = simpleobsws.IdentificationParameters() # 处理身份验证信息
ws = simpleobsws.WebSocketClient(url='ws://localhost:4455', identification_parameters=parameters)


# 切换场景的函数
async def switch_scene(scene_name):
    request = simpleobsws.Request('SetCurrentProgramScene', {'sceneName': scene_name})
    response = await ws.call(request) # 调用API请求
    if response.ok():
        print(f"成功切换到场景: {scene_name}")
    else:
        print(f"切换场景时出错: {response.responseData}")

# 截图函数
async def take_screenshot(scene_name):
    timestamp = time.strftime("%Y%m%d%H%M%S", time.localtime()) # 获取当前时间戳
    screenshot_path = f"/home/zengkai/Codes/obs/pics/{scene_name}_screenshot_{timestamp}.png" # 修改为带时间戳的保存路径
    request = simpleobsws.Request('SaveSourceScreenshot', {
    'sourceName': scene_name,
    'imageFormat': 'png',
    'imageFilePath': screenshot_path
    })
    response = await ws.call(request)
    if response.ok():
        print(f"截图成功，保存路径: {screenshot_path}")
    else:
        print(f"截图时出错: {response.responseData}")

# 主函数，连接 WebSocket 并抓取各个场景画面
async def main():
    await ws.connect() # 连接 WebSocket
    await ws.wait_until_identified() # 等待 OBS WebSocket 完全连接

    # scenes = ["s1", "s2", "s3"] # 替换为你实际的场景名称

    # while True:
    #     for scene in scenes:
    #         await switch_scene(scene) # 切换到指定场景
    #         time.sleep(0.5) 
    #         await take_screenshot(scene) # 截取当前场景的截图
            
    #         if scene == scenes[-1]:
    #             time.sleep(5) # 每 5 秒抓取一个场景的画面

    scene = "s-all"
    await switch_scene(scene)
    time.sleep(0.1) 
    while True:
        await take_screenshot(scene)
        time.sleep(1) 

    await ws.disconnect() # 断开 WebSocket 连接

# # 运行主函数
# asyncio.run(main())

# 获取事件循环并运行主函数
if __name__ == "__main__":
    try:
        loop = asyncio.get_event_loop() # 获取当前的事件循环
        loop.run_until_complete(main()) # 在事件循环中运行主函数\
    except KeyboardInterrupt:
        # loop.run_until_complete(ws.disconnect())
        # loop.close()
        print("程序已结束")
    except Exception as e:
        print(f"程序发生错误: {e}")
        # loop.run_until_complete(ws.disconnect())
        # loop.close()
        print("程序已结束")
    